import random
import sys
import itertools
import math
import bisect

# TODO
# * Generalize fuzzers that use prefixes
# * use seed

class Fuzzer(object):
    '''The "Abstract" fuzzer super class.'''
    def __init__(self, samples, seed=None, fuzzfactor=0.02):
        self.samples = samples
        self.seed = seed
        self.ff = fuzzfactor

        #Indicates whether the incrementation should be used or not
        self.incrementation = False

    def getPrefixes(self):
        # TODO
        return ['']

    def toBytes(self, data):
        if data == 0:
            return '\x00'

        bytes = ''
        while data > 0:
            bytes += chr(data & 255)
            data = data >> 8

        return bytes

    def setIncrementation(self):
        self.incrementation = True

    '''Incrementation feature should not be mandatory to use'''
    def increment(self, samples):
        if self.incrementation:
            self.samples.extend(samples)
        else:
            self.samples = samples

    def setFuzzFactor(self, ff):
        self.ff = ff

class Enumr(Fuzzer):
    """Enumr - the byte enumerator

    Enumr just starts from byte 0 and moves to byte 1 and so on. It doesn't need any input to function"""
    
    def __init__(self, samples, seed=None):
        super(Enumr, self).__init__(samples, seed)
        self.prefixes = self.getPrefixes()
        self.byte = -1

    def __iter__(self):
        return self

    def next(self):
        for prefix in self.prefixes:
            self.byte += 1
        return prefix + self.toBytes(self.byte)

class Flipr(Fuzzer):
    """ Flipr - the byte flipper

    Flipr randomly selects some places in on of the samples and flips
    one randomly chosen bit. The amount of flips depends on the fuzzfactor."""
    
    def __init__(self, samples, seed=None, ff=0.02):
        super(Flipr, self).__init__(samples, seed, ff)
        self.prefixes = self.getPrefixes()

    def __iter__(self):
        return self

    def next(self):

        sample = list(random.choice(self.samples))

        maxtimes = max(1, int(round((self.ff / 10) * len(sample) * 2)))
        times = random.randint(1, maxtimes)

        for i in range(times): 
            index = random.randrange(len(sample)) * 8
            pos = index / 8

            byteOffset = index % 8

            byte = ord(sample[pos])
            byte = byte ^ (1 << byteOffset)
            sample[pos] = chr(byte)

        return ''.join(sample)
                
class Rambo(Fuzzer):
    """ Rambo - the file mutator

    Rambo uses probably the most common fuzzing technique. It selects
    some random places in a randomly picked sample and removes,
    repeats, or replaces a byte. The amount of mutations depends on
    the fuzzfactor.
    """
    def __init__(self, samples, seed=None, ff=0.02):
        super(Rambo, self).__init__(samples, seed, ff)

    def __iter__(self):
        return self

    def next(self):
        sample = list(random.choice(self.samples))

        maxtimes = max(1, int(round((self.ff / 10) * len(sample) * 2)))
        times = random.randint(1, maxtimes)

        for i in range(times):
            index = random.randrange(len(sample))
            type = random.randrange(3)
            if type == 0:
                #remove a byte
                sample[index] = ''

            elif type == 1:
                #repeat a byte
                sample[index] = sample[index] * 2

            elif type == 2:
                #insert some random rata
                sample[index] = chr(random.randrange(256))

        return ''.join(sample)

class Range(Fuzzer):
    """Range - the random data generator

    Range generates an test case filled with pseudo-random content.
    """

    def __init__(self, samples, seed=None):
        super(Range, self).__init__(samples, seed)
        self.prefixes = self.getPrefixes()

    def __iter__(self):
        return self

    def next(self):
        # FIXME use some nice distribution
        times = random.randrange(10**3)
        return "".join([chr(random.randrange(256)) for x in range(times)])


class DistFuzz(Fuzzer):
    """ DiztFuzz - a byte distribution based fuzzer

    Distfuzz calculates the byte distribution of all samples and
    generates test cases so that more probable bytes are more probable
    in the test cases. It also, depending on the fuzzfactor, inserts
    some pseudo-random bytes in some places.
    """
    def __init__(self, samples, seed=None, ngramSize=1, ff=0.02):
        super(DistFuzz, self).__init__(samples, seed, ff)
        self.prefixes = self.getPrefixes()

        self.counts = dict()
        self.sizes = list()
        self.ngramSize = ngramSize

        #self.increment(samples)


    def increment(self, samples):
        # FIXME You don't have to calculate the differences from
        # scratch every time you increment. I was feeling lazy (and
        # stupid) so I didn't bother to do this sensibly.

        if not self.incrementation:
            self.samples = samples
        else:
            self.samples.extend(samples)

        for sample in self.samples:
            self.sizes.append(len(sample))

            for ngram, count in self.countGrams(sample).items():
                self.counts[ngram] = self.counts.get(ngram, 0) + count


        # We calculate the mean and standard deviation of file sizes. We
        # use this later to randomply pick file sizes
        #print 'sizes', self.sizes
        self.fmean = sum(self.sizes) / len(self.sizes) 
        self.fstddev = math.sqrt(sum([(x - self.fmean) ** 2 for x in self.sizes]) / len(self.sizes))
        self.distribution = dict()

        # We now match every occurences to the "propability mass" of the ngram
        # The following diagram tries to illusrate this
        # counts:
        #     a: 5, b: 10, c: 5
        # ->  a = 1-5, b = 6-16, c = 17-22
        # ->  5: a, 16: b, c: 22
        # new dict of propabilities --^

        first = True
        self.total_count = 0

        for ngram, count in self.counts.items():
            if first:
                count -= 1
                first = False
            self.total_count += count + 1
            self.distribution[self.total_count] = ngram

    def next(self):
        keys = self.distribution.keys()
        keys.sort()

        content = list()

        size = self.getSize()
        ifuzz = self.pickIndices(size)

        for j in range(size / self.ngramSize):
            if j in ifuzz:
                ngram = [chr(random.randrange(255)) for x in range(self.ngramSize)]
            else:
                r = random.randint(1, self.total_count)
                #we use the bisect module to see where r "hits" in
                #the propability mass
                order = bisect.bisect_left(keys, r)
                ngram = self.distribution[keys[order]]
            content.extend(ngram)
        return ''.join(content)

    def __iter__(self):
        return self
        
    #For handling n-grams from Jaochim Wudde's incmental
    def countGrams(self, data):
        n = self.ngramSize
        '''Counts ngrams from given string'''
        counts = dict()
        for ngram in itertools.izip(*[data[i:] for i in range(n)]):
            counts[ngram] = counts.get(ngram, 0) + 1
        return counts

    def pickIndices(self, size):
        '''Randomly chooses which indeces are fuzzed. Uses gaussian
        estimation for binomial distribution. Returns list of indices.'''
        n = float(size)
        mean = n * self.ff / 10
        stddev = n * self.ff / 10 * (1 - self.ff / 10)
        # we approximate te binomial distrobution with gaussian distribution.
        number = int(max(0, round(random.gauss(mean, stddev))))

        #Fix to prevent random.sample choosing more variables than it's able to
        if number > size: number = size
        
        indices = random.sample(xrange(size), number)
        return indices

    def getSize(self):
        '''returns a randomly chosen file size using gaussian distribution'''
        return max(1, int(round(random.gauss(self.fmean, self.fstddev))))


'''
This class is meant to be an interface between network hoover and fuzzer
modules defined here. Main purpose of this class is to separate there two
so that it would be easier to expand existing fuzzers without changing the
implementation of the whole hoover too much.
'''
class Launcher():
    
    def __init__(self, flist, fuzzer_opt):
        self.choice = {}
        self.fopt = {}

        fuzzers = ['Flipr','Enumr','Rambo','Range','DistFuzz']

        i = 0
        temp = []

        while i<len(fuzzer_opt):
            
            if fuzzer_opt[i] in fuzzers:

                temp = [fuzzer_opt[i]]
                i+=1

                while fuzzer_opt[i] not in fuzzers:
                    temp.append(fuzzer_opt[i])
                    if i==len(fuzzer_opt)-1: break
                    else: i+=1

                #This is where wanted values are parsed and rest of them are ignored
                if 'increment' in temp:
                    self.fopt[temp[0]+'-increment'] = True
                if 'fuzzfactor' in temp:
                    try:
                        self.fopt[temp[0]+'-fuzzfactor']=float(temp[temp.index('fuzzfactor')+1])
                    except ValueError:
                        print "Fuzzfactor value erronous for the fuzzer '{}'".format(temp[0])
                        self.fopt[temp[0]+'-fuzzfactor'] = 0.02
            else: i+=1

        '''
        Depending on given options will fuzzers be initialized in a correct way.
        Every fuzzer has to be initialized separately so that it'll easier to
        give different values to them.
        '''
        if 'Flipr' in flist:
            self.Flipr = Flipr([])
            self.choice.update({'Flipr':self.Flipr})
            if 'Flipr-fuzzfactor' in self.fopt:
                self.Flipr.setFuzzFactor(self.fopt['Flipr-fuzzfactor'])
            if 'Flipr-increment' in self.fopt:
                self.Flipr.setIncrementation()

        if 'Enumr' in flist:
            self.Enumr = Enumr([])
            self.choice.update({'Enumr':self.Enumr})

        if 'Rambo' in flist:
            self.Rambo = Rambo([])
            self.choice.update({'Rambo':self.Rambo})
            if 'Rambo-fuzzfactor' in self.fopt:
                self.Rambo.setFuzzFactor(self.fopt['Rambo-fuzzfactor'])
            if 'Rambo-increment' in self.fopt:
                self.Rambo.setIncrementation()

        if 'Range' in flist:
            self.Range = Range([])
            self.choice.update({'Range':self.Range})

        if 'DistFuzz' in flist:
            self.DistFuzz = DistFuzz([])
            self.choice.update({'DistFuzz':self.DistFuzz})
            if 'DistFuzz-fuzzfactor' in self.fopt:
                self.DistFuzz.setFuzzFactor(self.fopt['DistFuzz-fuzzfactor'])
            if 'DistFuzz-increment' in self.fopt:
                self.DistFuzz.setIncrementation()

    def dataFuzz(self, data, fuzzer):

        #TODO: is there a limitation to incrementation

        '''If fuzzer doesn't exist, return the same data back'''
        if fuzzer not in self.choice:
            return data

        self.choice[fuzzer].increment([data])
        return self.choice[fuzzer].next()
                
    
if __name__ == '__main__':
    #samples = [open(x).read() for x in sys.argv[1:]]
    samples = "Sample data!"
    for fuzzer in [
        Enumr(samples),
        Flipr(samples, ff=0.02),
        Rambo(samples, ff=0.02),
        Range(samples),
        DistFuzz(samples, ff=0.02)
        ]:
        print fuzzer
    
        i = 0
        for data in fuzzer:
            print repr(data)
            if i > 10:
                break
            i += 1
